package com.mobile.embedding;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * 中国移动的AI客服助手规则
 * 在应用启动后异步将客服助手规则文件转换为向量并存储到向量数据库
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class RuleEmbedding {

    private final VectorStore vectorStore;
    private final ExecutorService embeddingExecutor;

    @Value("classpath:rules.json")
    private Resource resource;

    private volatile boolean initialized = false;
    private volatile boolean initializing = false;

    /**
     * 应用启动时初始化方法
     */
    @PostConstruct
    public void init() {
        embeddingExecutor.submit(this::asyncInit);
    }

    /**
     * 异步初始化AI客服助手规则
     */
    private void asyncInit() {
        if (initializing || initialized) {
            return;
        }

        initializing = true;
        try {
            log.info("开始异步初始化AI客服助手规则");

            // 1. 创建文本读取器并且加载文件内容
            TextReader textReader = new TextReader(resource);
            textReader.getCustomMetadata().put("filename", "citys.txt");

            // 2. 将文件内容拆分为小块文档
            List<Document> documentList = textReader.get();
            TokenTextSplitter textSplitter = new TokenTextSplitter(200, 100, 5, 10000, false);
            List<Document> splitDocuments = textSplitter.split(documentList);

            // 3. 将处理后的文档向量化并存入到向量存储中
            vectorStore.add(splitDocuments);

            initialized = true;
            log.info("数据写入向量库成功，数据条数：{}", splitDocuments.size());
        } catch (Exception e) {
            log.error("异步初始化AI客服助手规则失败", e);
        } finally {
            initializing = false;
        }
    }

    /**
     * 检查是否已完成初始化
     * @return true表示已完成初始化
     */
    public boolean isInitialized() {
        return initialized;
    }

    /**
     * 检查是否正在初始化
     * @return true表示正在初始化
     */
    public boolean isInitializing() {
        return initializing;
    }

    /**
     * 等待初始化完成
     * @param timeoutMillis 超时时间(毫秒)
     * @throws InterruptedException 线程中断异常
     * @throws TimeoutException 超时异常
     */
    public void waitForInitialization(long timeoutMillis) throws InterruptedException, TimeoutException {
        long startTime = System.currentTimeMillis();
        while (initializing && !initialized) {
            if (System.currentTimeMillis() - startTime > timeoutMillis) {
                throw new TimeoutException("等待初始化超时");
            }
            Thread.sleep(100);
        }

        if (!initialized) {
            throw new IllegalStateException("初始化未完成");
        }
    }

    /**
     * 应用关闭时优雅关闭线程池
     */
    @PreDestroy
    public void destroy() {
        if (embeddingExecutor != null) {
            embeddingExecutor.shutdown();
            try {
                if (!embeddingExecutor.awaitTermination(60, TimeUnit.SECONDS)) {
                    embeddingExecutor.shutdownNow();
                }
            } catch (InterruptedException e) {
                embeddingExecutor.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
    }
}

